package org.zc.hive;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFUtils;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorConverter;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;

import static org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector.PrimitiveCategory.*;

/**
 * @program: hive-function
 * @description: org.zc.hive.UpperFirstChar
 * @author: zengchen
 * @create: 2021-04-21 10:15
 **/
public class UpperFirstChar extends GenericUDF {

    private transient PrimitiveObjectInspectorConverter.StringConverter stringConverter;
    private transient GenericUDFUtils.StringHelper returnHelper;

    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        // 1.检查参数个数
        if(arguments.length != 1) {
            throw new UDFArgumentException("必须有且仅有一个参数！");
        }

        // 2.检查参数类型
        // 2.1 不能是 LIST, MAP, STRUCT, UNION
        if(arguments[0].getCategory() != ObjectInspector.Category.PRIMITIVE){
            throw new UDFArgumentException("仅支持原始数据类型！");
        }

        // 获取到 UDF的参数监视器
        PrimitiveObjectInspector inIO = (PrimitiveObjectInspector) arguments[0];

        // 2.2 原始数据类型也很多，进一步过滤，排除
        // VOID, BOOLEAN, BYTE, SHORT, INT, LONG, FLOAT, DOUBLE, STRING, DATE, TIMESTAMP, TIMESTAMPLOCALTZ,
        // BINARY, DECIMAL, VARCHAR, CHAR, INTERVAL_YEAR_MONTH, INTERVAL_DAY_TIME, UNKNOWN;
        PrimitiveObjectInspector.PrimitiveCategory inCategory = inIO.getPrimitiveCategory();
//      STRING VARCHAR CHAR
        if(inCategory != CHAR && inCategory != VARCHAR && inCategory != STRING){
            throw new UDFArgumentException("Unexpected non-string type " + inCategory.name());
        }

        // 初始化 evaluate() 方法要使用到的变量
        stringConverter = new PrimitiveObjectInspectorConverter.StringConverter(inIO);
        returnHelper = new GenericUDFUtils.StringHelper(inCategory);

        ObjectInspector outIO = PrimitiveObjectInspectorFactory.writableStringObjectInspector;
        if(inCategory == CHAR){
            outIO = PrimitiveObjectInspectorFactory.getPrimitiveWritableObjectInspector(
                    TypeInfoFactory.getCharTypeInfo(GenericUDFUtils.StringHelper.getFixedStringSizeForType(inIO)));
        }
        if(inCategory == VARCHAR){
            outIO = PrimitiveObjectInspectorFactory.getPrimitiveWritableObjectInspector(
                    TypeInfoFactory.getVarcharTypeInfo(GenericUDFUtils.StringHelper.getFixedStringSizeForType(inIO)));
        }

        return outIO;
    }


    public Object evaluate(DeferredObject[] arguments) throws HiveException {
        String val = (String) stringConverter.convert(arguments[0].get());
        return returnHelper.setReturnValue(val.substring(0,1).toUpperCase() + val.substring(1));
    }

    /**
     * Get the String to be displayed in explain.
     *
     * @param children
     */
    public String getDisplayString(String[] children) {
        return "执行失败之后，已经不想说话了";
    }

}
